
#include <immintrin.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <time.h>
#include <string.h>

#include "utils.h"
#include "amx_bench.h"
#include "vnni_avx512_bench.h"

// general interface for microbenchmarking
typedef int (*benchamx_if)(int64_t num_iter, int8_t *A, int8_t *B, int32_t *C, size_t M, size_t N, size_t K);

typedef struct {
    const char* name;    // benchmark name
    const char* driver;  // hardware component to bench, amx/avx/...
    const char* level;   // memory hierarchy, reg/l1/l2/...
    const bool pack;     // pack or not
    benchamx_if func;    // function to do microbenchmarking
} benchamx_entry;

benchamx_entry benchmarks[] = {
    {"tmm_4ld4tdpb_with_dep", "amx", "reg", false, tmm_4ld4tdpb_with_dep},
    {"tmm_4tdpb4st_with_dep_1", "amx", "reg", false, tmm_4tdpb4st_with_dep_1},
    {"tmm_4tdpb4st_with_dep_2", "amx", "reg", false, tmm_4tdpb4st_with_dep_2},
    {"tmm_4tdpb4st_without_dep_1", "amx", "reg", false, tmm_4tdpb4st_without_dep_1},
    {"tmm_4tdpb4st_without_dep_2", "amx", "reg", false, tmm_4tdpb4st_without_dep_2},
    {"tmm_4ld4tdpb4st_with_dep_100", "amx", "reg", false, tmm_4ldtdpb4st_with_dep_100},
    {"l1", "amx", "l1", false, l1},
    {"l1_bw_load_store", "amx", "l1", false, l1_bw_load_store},
    {"l1_bw_load", "amx", "l1", false, l1_bw_load},
    {"l1_bw_store", "amx", "l1", false, l1_bw_store},
    {"l2_v1", "amx", "l2", false, l1},
    {"l2_bw_load_store_v1", "amx", "l2", false, l1_bw_load_store},
    {"l2_bw_load_v1", "amx", "l2", false, l1_bw_load},
    {"l2_bw_store_v1", "amx", "l2", false, l1_bw_store},
    {"l2_v2", "amx", "l2", false, l2_v2},
    {"l2_v2_load", "amx", "l2", false, l2_v2_load},
    {"l2_v2_with_pack", "amx", "l2", true, l2_v2_with_pack},
    {"l2_v2_with_pack_swpfa", "amx", "l2", true, l2_v2_with_pack_swpfa},
    {"l2_v2_with_pack_swpfc", "amx", "l2", true, l2_v2_with_pack_swpfc},
    {"l2_v2_with_pack_swpfb", "amx", "l2", true, l2_v2_with_pack_swpfb},
    {"l2_v2_load_with_pack", "amx", "l2", true, l2_v2_load_with_pack},
    {"vnni_avx512_zmm", "vnni_avx512", "reg", false, vnni_avx512_zmm},
    {NULL, NULL, NULL, false, NULL}
};

benchamx_entry* find_benchmark(const char* name) {
    for (int i = 0; benchmarks[i].name != NULL; ++i) {
        if (strcmp(benchmarks[i].name, name) == 0) {
            return &benchmarks[i];
        }
    }

    return NULL;
}

typedef struct {
    double (*get_peak_gflops)();
    void (*b_layout_transform)(uint8_t*, uint8_t*, size_t, size_t);
    void (*stats)(const char*, const size_t, const size_t, const size_t, const double, const double);
} benchamx_utils;

int main(int argc, char** argv) {
    if (argc != 8) {
        printf("Usage: %s name nwarmups niters M N K loop_order\n", argv[0]);
        exit(1);
    }

    const char* name = argv[1];
    int num_warmups = atoi(argv[2]);
    int num_iters = atoi(argv[3]);
    size_t M = (size_t)atoi(argv[4]);
    size_t N = (size_t)atoi(argv[5]);
    size_t K = (size_t)atoi(argv[6]);
    const char* loop_order = argv[7];

    int8_t *A = (int8_t *)_mm_malloc(M * K * sizeof(int8_t), 64);
    int8_t *B = (int8_t *)_mm_malloc(K * N * sizeof(int8_t), 64);
    int8_t *B_trans = (int8_t *)_mm_malloc(K * N * sizeof(int8_t), 64);
    int32_t *C = (int32_t *)_mm_malloc(M * N * sizeof(int32_t), 64);

    benchamx_entry* bench = find_benchmark(name);
    benchamx_utils utils;

    stats(M, N, K);

    if (strcmp(bench->driver, "amx") == 0) {
        utils.get_peak_gflops = get_amx_peak_gflops;
        if (bench->pack) {
            utils.b_layout_transform = amx_b_layout_transform_with_pack;
        } else {
            utils.b_layout_transform = amx_b_layout_transform;
        }
    } else if (strcmp(bench->driver, "vnni_avx512") == 0) {
        utils.get_peak_gflops = get_vnni_avx512_peak_gflops;
        utils.b_layout_transform = vnni_avx512_b_layout_transform;
    } else {
        printf("unsupported driver! exiting...\n");
        exit(1);
    }

    if (strcmp(bench->level, "l1") == 0) {
        utils.stats = l1_stats;
    } else if (strcmp(bench->level, "l2") == 0) {
        utils.stats = l2_stats;
    } else {
        utils.stats = NULL;
    }


    init_matrices(A, B, C, M, N, K);

    utils.b_layout_transform(B, B_trans, K, N);

    if (strcmp(bench->driver, "amx") == 0) {
        // Request permission to linux kernel to run AMX
        if (!set_tiledata_use())
            exit(-1);
        // Load tile configuration
        __tilecfg tile_data = {0};
        init_tile_config(&tile_data);
    }

    // correctness test
    bench->func(1, A, B_trans, C, M, N, K);
    // print_matrix_to_file(A, M, K, 1, "A.txt");
    // print_matrix_to_file(B, K, N, 1, "B.txt");
    // print_matrix_to_file(C, M, N, 4, "C.txt");
    verify_matrix_multiplication(A, B, C, M, N, K, bench->pack);

    // warmup
    int ret = bench->func(num_warmups, A, B_trans, C, M, N, K);

    // run test
    struct timespec start, end;
    clock_gettime(CLOCK_MONOTONIC, &start);

    #if ENABLE_TRACE == 1
    __PINBALL_TRACE_START__();
    #endif

    ret = bench->func(num_iters, A, B_trans, C, M, N, K);

    #if ENABLE_TRACE == 1
    __PINBALL_TRACE_STOP__();
    #endif

    clock_gettime(CLOCK_MONOTONIC, &end);
 
    // Calculate elapsed time in seconds
    double elapsed = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
    double ops = 0.0;
    if (strstr(name, "tmm") != NULL) {
        ops = 2.0 * TILE_M * TILE_N * TILE_K * 4 * num_iters;
    }
    else {
        ops = 2.0 * M * N * K * num_iters;
    }

    // utils.stats(loop_order, M, N, K, ops/num_iters, elapsed/num_iters);

    double gflops = (ops / elapsed) / 1e9;
    printf("Performance:     [%.2f] GFLOPS\n", gflops);
    printf("MAC_Utilization: %.2f%%\n", (gflops / utils.get_peak_gflops()) * 100.0); 
 
    return 0;
 }